/**
  ******************************************************************************
  * @file           : studio_tcp_server.h
  * @author         : wangyingjie
  * @brief          : socket tcp server 类封装 基于 boost::asio
  * @attention      : None
  * @date           : 2025/7/1
  ******************************************************************************
  */

#ifndef STUDIO_TCP_SERVER_HPP
#define STUDIO_TCP_SERVER_HPP

#include "studio_macros.h"


#include <boost/asio.hpp>
#include <boost/array.hpp>
#include <boost/system/error_code.hpp>


namespace asio = boost::asio;
using tcp = asio::ip::tcp;
using boost::system::error_code;

class studio_tcp_server;

// 单个客户端连接会话类
class tcp_session : public std::enable_shared_from_this<tcp_session>
{
public:
    using message_handler = std::function<void(const std::string&)>;
    using disconnect_handler = std::function<void(std::shared_ptr<tcp_session>)>;

    explicit tcp_session(tcp::socket socket, studio_tcp_server* server)
        : socket_(std::move(socket)),
          server_(server),
          connected_(true)
    {
    }
    ~tcp_session()
    {
        close();
    }

    void start()
    {
        do_read();
    }

    void send(const std::string& message)
    {
        if (!connected_.load())
            return;

        bool expected = false;
        if (write_in_progress_.compare_exchange_strong(expected, true))
        {
            asio::async_write(socket_,
                              asio::buffer(message),
                              [this](const error_code& ec, std::size_t /*bytes*/) {
                                  write_in_progress_.store(false);
                                  if (ec)
                                  {
                                      handle_disconnect();
                                  }
                              });
        }
    }

    tcp::endpoint get_remote_endpoint()
    {
        return socket_.remote_endpoint(ec_);
    }

    void close()
    {
        boost::system::error_code ec;
        if (socket_.is_open())
        {
            socket_.close(ec);
        }
    }

private:
    void do_read()
    {
        socket_.async_read_some(asio::buffer(buffer_),
                                [this](const error_code& ec, std::size_t bytes_transferred) {
                                    if (ec)
                                    {
                                        handle_disconnect();
                                        return;
                                    }

                                    std::string data(buffer_.data(), bytes_transferred);
                                    if (msg_handler_)
                                    {
                                        msg_handler_(data);
                                    }

                                    do_read(); // 继续读取
                                });
    }

    void handle_disconnect()
    {
        connected_.store(false);
        if (disconnect_handler_)
        {
            disconnect_handler_(shared_from_this());
        }
    }

public:
    message_handler msg_handler_;
    disconnect_handler disconnect_handler_;

private:
    tcp::socket socket_;
    boost::array<char, 4096> buffer_;
    std::atomic<bool> connected_;
    std::atomic<bool> write_in_progress_{false};
    error_code ec_;
    studio_tcp_server* server_;
};

// TCP 服务端类
class studio_tcp_server
{
public:
    using message_handler = std::function<void(const std::string&, std::shared_ptr<tcp_session>)>;
    using disconnect_handler = std::function<void(std::shared_ptr<tcp_session>)>;

    studio_tcp_server(const std::string& ip, const std::string& port, int reconnect_interval = 3)
        : io_(),
          acceptor_(io_, tcp::endpoint(boost::asio::ip::make_address(ip), std::stoi(port))),
          reconnect_timer_(io_),
          reconnect_interval_(reconnect_interval),
          has_client_connected_(false)
    {
    }

    ~studio_tcp_server()
    {
        stop();
    }

    void set_message_handler(message_handler handler)
    {
        msg_handler_ = handler;
    }

    void set_disconnect_handler(disconnect_handler handler)
    {
        disconnect_handler_ = handler;
    }

    void start()
    {
        std::cout << "Starting TCP server on " << acceptor_.local_endpoint() << std::endl;
        do_accept();
        worker_thread_ = std::make_unique<std::thread>([this]() {
            io_.run();
        });
    }

    void stop()
    {
        //if (stopped_)
        //    return;
        //stopped_ = true;

        //io_.post([this]() {
        //    error_code ec;
        //    acceptor_.close(ec);
        //    reconnect_timer_.cancel(ec);

        //    for (auto& session : sessions_) {
        //        session.second->close();  // 关闭 socket
        //    }
        //    sessions_.clear();  // 清除 map 中的所有 shared_ptr，触发析构
        //    sessions_.clear();
        //});

        //if (worker_thread_ && worker_thread_->joinable())
        //{
        //    worker_thread_->join();
        //}
    }

    void broadcast(const std::string& message)
    {
        for (auto& session : sessions_)
        {
            session.second->send(message);
        }
    }

    void send_to_client(const std::string& client_id, const std::string& message)
    {
        std::lock_guard<std::mutex> lock(sessions_mutex_);
        auto it = sessions_.find(client_id);
        if (it != sessions_.end())
        {
            it->second->send(message);
        }
    }

private:
    void do_accept()
    {
        if (stopped_)
            return;

        acceptor_.async_accept(
            [this](error_code ec, tcp::socket socket) {
                if (!ec)
                {
                    auto session = std::make_shared<tcp_session>(std::move(socket), this);
                    std::string client_id = session->get_remote_endpoint().address().to_string() + ":" +
                                            std::to_string(session->get_remote_endpoint().port());

                    session->msg_handler_ = [this, client_id, session](const std::string& msg) {
                        if (msg_handler_)
                        {
                            msg_handler_(msg, session);
                        }
                    };

                    session->disconnect_handler_ = [this, client_id, session](std::shared_ptr<tcp_session>) {
                        std::lock_guard<std::mutex> lock(sessions_mutex_);
                        sessions_.erase(client_id);
                        check_and_reconnect();

                        if (disconnect_handler_)
                        {
                            disconnect_handler_(session);
                        }
                    };

                    {
                        std::lock_guard<std::mutex> lock(sessions_mutex_);
                        sessions_[client_id] = session;
                    }

                    session->start();

                    has_client_connected_ = true; // 标记已有客户端连接过
                }
                else
                {
                    std::cerr << "Accept failed: " << ec.message() << ". Retrying...\n";
                    check_and_reconnect();
                    return;
                }

                do_accept(); // 继续监听新连接
            });
    }

    void check_and_reconnect()
    {
        //if (stopped_)
        //    return;

        //if (sessions_.empty() && has_client_connected_)
        //{
        //    std::cout << "All clients disconnected. Waiting to restart listener...\n";
        //    reconnect_timer_.expires_from_now(std::chrono::seconds(reconnect_interval_));
        //    reconnect_timer_.async_wait([this](const error_code&) {
        //        io_.post([this]() {
        //            restart_accept();
        //        });
        //    });
        //}
    }

    void restart_accept()
    {
        if (stopped_)
            return;

        error_code ec;
        acceptor_.close(ec);
        acceptor_.open(tcp::v4());
        acceptor_.bind(tcp::endpoint(boost::asio::ip::make_address("0.0.0.0"), acceptor_.local_endpoint().port()));
        acceptor_.listen();

        do_accept();
    }

private:
    asio::io_context io_;
    tcp::acceptor acceptor_;
    std::unique_ptr<std::thread> worker_thread_;
    std::unordered_map<std::string, std::shared_ptr<tcp_session>> sessions_;
    std::mutex sessions_mutex_;

    asio::steady_timer reconnect_timer_;
    int reconnect_interval_;

    message_handler msg_handler_;
    disconnect_handler disconnect_handler_;

    std::atomic<bool> has_client_connected_;
    std::atomic<bool> stopped_{false};
};


#endif // STUDIO_TCP_SERVER_HPP